#include "main/utils.h"
#include "lib/binary_tree.h"
using namespace std;

using TreeNode = BinaryTreeNode;

class Solution {
public:
    int pathSum(TreeNode* root, int sum) {
      prefix[0] = 1;
      return dfs(root, 0, sum);
    }

    int dfs(TreeNode* root, long long cur, int sum) {
      if (!root)
        return 0;

      int ret = 0;
      cur += root->value;
      if (prefix.count(cur - sum)) {
        ret = prefix[cur - sum];
      }
      ++prefix[cur];
      ret += dfs(root->left, cur, sum);
      ret += dfs(root->right, cur, sum);
      --prefix[cur];

      return ret;
    }

    int pathSum1(TreeNode* root, int sum) {
      if (!root)
        return 0;
      int ret = rootSum(root, sum);
      ret += pathSum1(root->left, sum);
      ret += pathSum1(root->right, sum);
      return ret;
    }

    int rootSum(TreeNode* root, int val) {
      if (!root)
        return 0;
      int ret = 0;
      if (root->value == val) {
        ++ret;
      }
      ret += rootSum(root->left, val - root->value);
      ret += rootSum(root->right, val - root->value);
      return ret;
    }

private:
    map<long long, int> prefix;
};

int main() {
//               5
//              / \
//             4   8
//            /   / \
//           11  13  4
//          /  \    / \
//         7    2  5   1
  auto node5 = CreateBinaryTreeNode(5);
  auto node4 = CreateBinaryTreeNode(4);
  auto node8 = CreateBinaryTreeNode(8);
  auto node11 = CreateBinaryTreeNode(11);
  auto node13 = CreateBinaryTreeNode(13);
  auto node4_1 = CreateBinaryTreeNode(4);
  auto node7 = CreateBinaryTreeNode(7);
  auto node2 = CreateBinaryTreeNode(2);
  auto node5_1 = CreateBinaryTreeNode(5);
  auto node1 = CreateBinaryTreeNode(1);

  ConnectNode(node5, node4, node8);
  ConnectNode(node4, node11, nullptr);
  ConnectNode(node8, node13, node4_1);
  ConnectNode(node11, node7, node2);
  ConnectNode(node4_1, node5_1, node1);

  PrintBinaryTree(node5);

  Solution su;
  cout << su.pathSum(node5, 22) << endl;
  cout << su.pathSum1(node5, 22) << endl;

  return 0;
}
